{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SzKwuqYESWwm"
   },
   "source": [
    "##### Copyright 2021 The Cirq Developers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "cellView": "form",
    "id": "4yPUsdJxSXFq"
   },
   "outputs": [],
   "source": [
    "# @title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "# https://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zC1qlUJoSXhm"
   },
   "source": [
    "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://quantumai.google/cirq/noise/qcvv/parallel_xeb\"><img src=\"https://quantumai.google/site-assets/images/buttons/quantumai_logo_1x.png\" />View on QuantumAI</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/quantumlib/Cirq/blob/main/docs/noise/qcvv/parallel_xeb.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/colab_logo_1x.png\" />Run in Google Colab</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a target=\"_blank\" href=\"https://github.com/quantumlib/Cirq/blob/main/docs/noise/qcvv/parallel_xeb.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/github_logo_1x.png\" />View source on GitHub</a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://storage.googleapis.com/tensorflow_docs/Cirq/docs/noise/qcvv/parallel_xeb.ipynb\"><img src=\"https://quantumai.google/site-assets/images/buttons/download_icon_1x.png\" />Download notebook</a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bd9529db1c0b"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    import cirq\n",
    "except ImportError:\n",
    "    print(\"installing cirq...\")\n",
    "    !pip install --quiet cirq\n",
    "    print(\"installed cirq.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dd95be2a71eb"
   },
   "source": [
    "# Parallel XEB\n",
    "This notebook demonstrates how to use the functionality in `cirq.experiments` to run parallel XEB end-to-end. \"Parallel\" means we characterize multiple pairs simultaneously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "046b07823210"
   },
   "outputs": [],
   "source": [
    "import cirq\n",
    "import numpy as np\n",
    "\n",
    "%matplotlib inline\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8bbaafb0ab88"
   },
   "source": [
    "# Parallel XEB with library functions\n",
    "The entire XEB workflow can be run by calling `cirq.experiments.parallel_two_qubit_xeb` and the combined single-qubit randomized benchmarking (RB) and XEB workflows can be run by calling `cirq.experiments.run_rb_and_xeb`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "c62af7442068"
   },
   "outputs": [],
   "source": [
    "# Simulation\n",
    "qubits = cirq.GridQubit.rect(3, 2, 4, 3)\n",
    "result = cirq.experiments.parallel_two_qubit_xeb(\n",
    "    sampler=cirq.DensityMatrixSimulator(\n",
    "        noise=cirq.depolarize(5e-3), dtype=np.complex128\n",
    "    ),  # Any simulator or a ProcessorSampler.\n",
    "    qubits=qubits,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0d3bb3069d30"
   },
   "outputs": [],
   "source": [
    "# The returned result is an instance of the `TwoQubitXEBResult` class which provides visualization methods like\n",
    "result.plot_heatmap()\n",
    "# plot the heatmap of XEB errors\n",
    "result.plot_fitted_exponential(*qubits[:2])\n",
    "# plot the fitted model of xeb error of a qubit pair.\n",
    "result.plot_histogram();  # plot a histogram of all xeb errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a3a8a3603e46"
   },
   "outputs": [],
   "source": [
    "# `TwoQubitXEBResult` also has methods to retrieve errors.\n",
    "print('pauli errors:', result.pauli_error())\n",
    "print('xeb errors:', result.xeb_error(*qubits[:2]))\n",
    "print('xeb fidelity:', result.xeb_fidelity(*qubits[:2]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "041c42d258ce"
   },
   "source": [
    "The `run_rb_and_xeb` method returns an object of type [InferredXEBResult](https://github.com/quantumlib/Cirq/blob/bc766606b94744f80da435c522d16a34529ae671/cirq-core/cirq/experiments/two_qubit_xeb.py#L188C7-L188C24) which is like [TwoQubitXEBResult](https://github.com/quantumlib/Cirq/blob/bc766606b94744f80da435c522d16a34529ae671/cirq-core/cirq/experiments/two_qubit_xeb.py#L56) except that it removes the single-qubit errors obtained from the single-qubit randomized benchmarking (RB) experiment to isolate the error from the two qubit gate."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ace31cc4d258"
   },
   "source": [
    "# Step by step XEB\n",
    "The rest of this notebook explains how the `parallel_two_qubit_xeb` works internally. Note that the notebook uses `SQRT_ISWAP` as the entangling gate while `parallel_two_qubit_xeb` and `run_rb_and_xeb` default to `CZ`.\n",
    "\n",
    "## Set up Random Circuits\n",
    "\n",
    "We create a library of 10 random, two-qubit `circuits` using the sqrt(ISWAP) gate. These library circuits will be mixed-and-matched among all the pairs on the device we aim to characterize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6fa8a0266b7a"
   },
   "outputs": [],
   "source": [
    "from cirq.experiments import random_quantum_circuit_generation as rqcg\n",
    "\n",
    "circuit_library = rqcg.generate_library_of_2q_circuits(\n",
    "    n_library_circuits=20, two_qubit_gate=cirq.ISWAP**0.5, random_state=52\n",
    ")\n",
    "print(len(circuit_library))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "c7c044ec12ac"
   },
   "outputs": [],
   "source": [
    "# We will truncate to these lengths\n",
    "max_depth = 100\n",
    "cycle_depths = np.arange(3, max_depth, 20)\n",
    "cycle_depths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6fbec4054504"
   },
   "source": [
    "## Determine the device topology\n",
    "\n",
    "We will run on all pairs from a given device topology. Below, you can supply a `device_name` if you're authenticated to run on Google QCS. In that case, we will get the device object from the cloud endpoint and turn it into a graph of qubits. Otherwise, we mock a device graph by allocating arbitrary `cirq.GridQubit`s to turn into a graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "af1e12169e88"
   },
   "outputs": [],
   "source": [
    "device_name = None  # change me!\n",
    "\n",
    "import cirq.contrib.routing as ccr\n",
    "import networkx as nx\n",
    "\n",
    "if device_name is None:\n",
    "    qubits = cirq.GridQubit.rect(3, 2, 4, 3)\n",
    "    # Delete one qubit from the rectangular arangement to\n",
    "    # 1) make it irregular 2) simplify simulation.\n",
    "    qubits = qubits[:-1]\n",
    "    sampler = cirq.DensityMatrixSimulator(noise=cirq.depolarize(5e-3))\n",
    "    graph = ccr.gridqubits_to_graph_device(qubits)\n",
    "else:\n",
    "    import cirq_google as cg\n",
    "\n",
    "    sampler = cg.get_engine_sampler(device_name, gate_set_name='sqrt_iswap')\n",
    "    device = cg.get_engine_device(device_name)\n",
    "    qubits = sorted(device.qubits)\n",
    "    graph = ccr.gridqubits_to_graph_device(device.qubits)\n",
    "\n",
    "\n",
    "pos = {q: (q.row, q.col) for q in qubits}\n",
    "nx.draw_networkx(graph, pos=pos)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5d6f40acc7bf"
   },
   "source": [
    "## Set up our combinations\n",
    "We take the library of two-qubit circuits in `circuit_library` and mix-and-match to sampled in parallel.\n",
    "\n",
    "We will pass `combs_by_layer` and `circuit_library` to the sampling function which will \"zip\" the circuits according to these combinations. The outer list corresponds to the four `cirq.GridInteractionLayer`s (one of four for the degree-four GridQubit-implied graph). The inner `combinations` matrix is a `(n_combinations, n_pairs)` ndarray of integers which index into the circuit library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fbaf9c5c5335"
   },
   "outputs": [],
   "source": [
    "combs_by_layer = rqcg.get_random_combinations_for_device(\n",
    "    n_library_circuits=len(circuit_library), n_combinations=10, device_graph=graph, random_state=53\n",
    ")\n",
    "combs_by_layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "d82866ff2269"
   },
   "source": [
    "### Visualize\n",
    "Here, we draw the four layers' active pairs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7ea374d2d337"
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 2, figsize=(9, 6))\n",
    "for comb_layer, ax in zip(combs_by_layer, axes.reshape(-1)):\n",
    "    active_qubits = np.array(comb_layer.pairs).reshape(-1)\n",
    "    colors = ['red' if q in active_qubits else 'blue' for q in graph.nodes]\n",
    "    nx.draw_networkx(graph, pos=pos, node_color=colors, ax=ax)\n",
    "    nx.draw_networkx_edges(\n",
    "        graph, pos=pos, edgelist=comb_layer.pairs, width=3, edge_color='red', ax=ax\n",
    "    )\n",
    "\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "423f56a942f6"
   },
   "source": [
    "## Take Data\n",
    "\n",
    "The following call will execute the zipped circuits and sample bitstrings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "aab08193f046"
   },
   "outputs": [],
   "source": [
    "from cirq.experiments.xeb_sampling import sample_2q_xeb_circuits\n",
    "\n",
    "sampled_df = sample_2q_xeb_circuits(\n",
    "    sampler=sampler,\n",
    "    circuits=circuit_library,\n",
    "    cycle_depths=cycle_depths,\n",
    "    combinations_by_layer=combs_by_layer,\n",
    "    shuffle=np.random.RandomState(52),\n",
    "    repetitions=10_000,\n",
    ")\n",
    "sampled_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "16b5a4a86384"
   },
   "source": [
    "## Benchmark Fidelities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8dd3ed7024ef"
   },
   "outputs": [],
   "source": [
    "from cirq.experiments.xeb_fitting import benchmark_2q_xeb_fidelities\n",
    "\n",
    "fids = benchmark_2q_xeb_fidelities(\n",
    "    sampled_df=sampled_df, circuits=circuit_library, cycle_depths=cycle_depths\n",
    ")\n",
    "fids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "05748707cacc"
   },
   "outputs": [],
   "source": [
    "from cirq.experiments.xeb_fitting import fit_exponential_decays, exponential_decay\n",
    "\n",
    "fidelities = fit_exponential_decays(fids)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WJbgW4r1Guuf"
   },
   "outputs": [],
   "source": [
    "heatmap_data = {}\n",
    "\n",
    "for (_, _, pair), fidelity in fidelities.layer_fid.items():\n",
    "    heatmap_data[pair] = 1.0 - fidelity\n",
    "\n",
    "cirq.TwoQubitInteractionHeatmap(heatmap_data).plot();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a12f73217150"
   },
   "outputs": [],
   "source": [
    "for i, record in fidelities.iterrows():\n",
    "    plt.axhline(1, color='grey', ls='--')\n",
    "    plt.plot(record['cycle_depths'], record['fidelities'], 'o')\n",
    "    xx = np.linspace(0, np.max(record['cycle_depths']))\n",
    "    plt.plot(xx, exponential_decay(xx, a=record['a'], layer_fid=record['layer_fid']))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8c08c9ab8109"
   },
   "outputs": [],
   "source": [
    "import seaborn as sns\n",
    "\n",
    "# Give each pair its own color\n",
    "colors = sns.cubehelix_palette(n_colors=graph.number_of_edges())\n",
    "colors = dict(zip(graph.edges, colors))\n",
    "\n",
    "# Exponential reference\n",
    "xx = np.linspace(0, fids['cycle_depth'].max())\n",
    "plt.plot(xx, (1 - 5e-3) ** (4 * xx), label=r'Exponential Reference', color='black')\n",
    "\n",
    "\n",
    "# Plot each pair\n",
    "def _p(fids):\n",
    "    q0, q1 = fids.name\n",
    "    plt.plot(\n",
    "        fids['cycle_depth'],\n",
    "        fids['fidelity'],\n",
    "        'o-',\n",
    "        label=f'{q0}-{q1}',\n",
    "        color=colors[fids.name],\n",
    "        alpha=0.5,\n",
    "    )\n",
    "\n",
    "\n",
    "fids.groupby('pair').apply(_p)\n",
    "\n",
    "plt.ylabel('Circuit fidelity')\n",
    "plt.xlabel('Cycle Depth $d$')\n",
    "plt.legend(loc='best')\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6025a292d19b"
   },
   "source": [
    "## Optimize `PhasedFSimGate` parameters\n",
    "\n",
    "We know what circuits we requested, and in this simulated example, we know what coherent error has happened. But in a real experiment, there is likely unknown coherent error that you would like to characterize. Therefore, we make the five angles in `PhasedFSimGate` free parameters and use a classical optimizer to find which set of parameters best describes the data we collected from the noisy simulator (or device, if this was a real experiment)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "baff45b4ad70"
   },
   "outputs": [],
   "source": [
    "import multiprocessing\n",
    "\n",
    "pool = multiprocessing.get_context('spawn').Pool()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a3aaf1dfce4e"
   },
   "outputs": [],
   "source": [
    "from cirq.experiments.xeb_fitting import (\n",
    "    parameterize_circuit,\n",
    "    characterize_phased_fsim_parameters_with_xeb_by_pair,\n",
    "    SqrtISwapXEBOptions,\n",
    ")\n",
    "\n",
    "# Set which angles we want to characterize (all)\n",
    "options = SqrtISwapXEBOptions(\n",
    "    characterize_theta=True,\n",
    "    characterize_zeta=True,\n",
    "    characterize_chi=True,\n",
    "    characterize_gamma=True,\n",
    "    characterize_phi=True,\n",
    ")\n",
    "# Parameterize the sqrt(iswap)s in our circuit library\n",
    "pcircuits = [parameterize_circuit(circuit, options) for circuit in circuit_library]\n",
    "\n",
    "# Run the characterization loop\n",
    "characterization_result = characterize_phased_fsim_parameters_with_xeb_by_pair(\n",
    "    sampled_df,\n",
    "    pcircuits,\n",
    "    cycle_depths,\n",
    "    options,\n",
    "    pool=pool,\n",
    "    # ease tolerance so it converges faster:\n",
    "    fatol=1e-2,\n",
    "    xatol=1e-2,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "c3961ac4e6c7"
   },
   "outputs": [],
   "source": [
    "characterization_result.final_params"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8fea4a364c79"
   },
   "outputs": [],
   "source": [
    "characterization_result.fidelities_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "86da92e214d6"
   },
   "outputs": [],
   "source": [
    "from cirq.experiments.xeb_fitting import before_and_after_characterization\n",
    "\n",
    "before_after_df = before_and_after_characterization(fids, characterization_result)\n",
    "before_after_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "583abc4f0c47"
   },
   "outputs": [],
   "source": [
    "for i, row in before_after_df.iterrows():\n",
    "    plt.axhline(1, color='grey', ls='--')\n",
    "    plt.plot(row['cycle_depths_0'], row['fidelities_0'], '*', color='red')\n",
    "    plt.plot(row['cycle_depths_c'], row['fidelities_c'], 'o', color='blue')\n",
    "\n",
    "    xx = np.linspace(0, np.max(row['cycle_depths_0']))\n",
    "    plt.plot(\n",
    "        xx,\n",
    "        exponential_decay(xx, a=row['a_0'], layer_fid=row['layer_fid_0']),\n",
    "        color='red',\n",
    "        label=f'f_0 = {row[\"layer_fid_0\"]:.3f}',\n",
    "    )\n",
    "    plt.plot(\n",
    "        xx,\n",
    "        exponential_decay(xx, a=row['a_c'], layer_fid=row['layer_fid_c']),\n",
    "        color='blue',\n",
    "        label=f'f_c = {row[\"layer_fid_c\"]:.3f}',\n",
    "    )\n",
    "\n",
    "    plt.xlabel('Cycle Depth')\n",
    "    plt.ylabel('Fidelity')\n",
    "    plt.legend(loc='best')\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "parallel_xeb.ipynb",
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
